Skip to content

Add metadata to dataset config files#420

Merged
jlamypoirier merged 34 commits intomainfrom
jlp_dataset_metadata
Dec 22, 2025
Merged

Add metadata to dataset config files#420
jlamypoirier merged 34 commits intomainfrom
jlp_dataset_metadata

Conversation

@jlamypoirier
Copy link
Collaborator

@jlamypoirier jlamypoirier commented Dec 12, 2025

✨ Description

Gather all available metadata from readers during dataset preparation and add it to the yaml config file. Works with both blending and splitting. Example:

$ cat /tmp/fast_llm_tests/common/dataset/dataset_with_image_patches_55_132/fast_llm_config.yaml 
config:
  path: shard_0_0.fast_llm_dataset
  type: memmap
metadata:
  image_patches:
    data_type: uint8
    num_documents: 1000
    num_patch_groups: 1355
    num_patches: 10707
    num_pixels: 513936
    patch_shape:
    - 3
    - 4
    - 4
  num_tokens: 59145
  tokens:
    data_type: int32
    num_documents: 1000
    num_tokens: 59145

The file format is different from before (config is now inside of a dict instead of at top level), but the old format can still be loaded.

@jlamypoirier
Fix rotary 2d
2ab1825
@jlamypoirier
stuff
8305dd5
@jlamypoirier
stuff
b6e38b8
@jlamypoirier
Merge branch 'main' into jlp/consistent_preprocessing
72f915d
@jlamypoirier
stuff
350fb3d
@jlamypoirier
fix
d27a815
@jlamypoirier
Merge branch 'main' into jlp/consistent_preprocessing
72f3a31
@jlamypoirier
fixes
5ab6cd0
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp/consistent_prepro…
1e74469 
…cessing
@jlamypoirier
stuff
6454db4
@jlamypoirier
cleanup
916af7a
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp/consistent_prepro…
355af7c 
…cessing
@jlamypoirier
Merge branch 'jlp/consistent_preprocessing' into jlp/varlen_tweaks
8f6841e
@jlamypoirier
cleanup
bd7a8e6
@jlamypoirier
fix
660fecc
@jlamypoirier
fixes
db93bb5
@jlamypoirier
Merge branch 'jlp/consistent_preprocessing' into jlp/varlen_tweaks
a3fa577
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp/varlen_tweaks
a1c0ade
@jlamypoirier
misc
96ce759
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp/varlen_tweaks
e23ea04
@jlamypoirier
stuff
e5fe8b2
@jlamypoirier
fixes
68f457b
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp/varlen_tweaks
fa668fa
@jlamypoirier
Remove mamba and discrete mamba 2
f7c5d1b
@jlamypoirier
fix
31d856d
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp_remove_mamba
e74d30d
@jlamypoirier
fixes
75ad78a
@jlamypoirier
Add metadata to dataset config files
30e0419
@jlamypoirier
fix
9fae16e
@jlamypoirier
Merge remote-tracking branch 'origin/main' into jlp_remove_mamba
2b6527a
Base automatically changed from jlp_remove_mamba to main December 13, 2025 06:21
@tscholak
Copy link
Collaborator

tscholak commented Dec 13, 2025

This looks very useful, thanks @jlamypoirier
@oleksost @RaymondLi0 is this close to what you had in mind?

@oleksost
Copy link
Contributor

oleksost commented Dec 15, 2025 •
edited
Loading

yes, this looks useful.
So the fast_llm_config.yaml will contain such entry per shard? Wouldn't it be more useful to have a dataset view isntead of per shard view?

@jlamypoirier
Copy link
Collaborator Author

jlamypoirier commented Dec 15, 2025

yes, this looks useful. So the fast_llm_config.yaml will contain such entry per shard? Wouldn't it be more useful to have a dataset view isntead of per shard view?

This is already global, the example just happens to have only one shard.

oleksost
oleksost approved these changes Dec 22, 2025
View reviewed changes
Copy link
Contributor

@oleksost oleksost left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are some merge conflicts, apart from that LHTM!

@jlamypoirier jlamypoirier merged commit 55c6d0e into main Dec 22, 2025
2 checks passed
@jlamypoirier jlamypoirier deleted the jlp_dataset_metadata branch December 22, 2025 17:51
tscholak
tscholak reviewed Dec 24, 2025
View reviewed changes
fast_llm/data/dataset/gpt/config.py
assert self.path.is_file(), f"File {self.path} does not exist."
return SampledDatasetConfig[SampleType].from_dict(self._convert_paths(yaml.safe_load(self.path.open("r"))))
config = yaml.safe_load(self.path.open("r"))
Assert.eq(config.keys(), {"config", "metadata"})
Copy link
Collaborator

@tscholak tscholak Dec 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jlamypoirier this is causing crashes now. remove?

Suggested change
Assert.eq(config.keys(), {"config", "metadata"})

Copy link
Collaborator Author

@jlamypoirier jlamypoirier Dec 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I must have left this there by accident

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@tscholak tscholak tscholak left review comments

@oleksost oleksost oleksost approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@jlamypoirier @tscholak @oleksost